import cv2
import mediapipe as mp
import numpy as np
import time
from sklearn.tree import DecisionTreeClassifier
import pickle
import os

# 初始化手部检测模型
mp_hands = mp.solutions.hands
hands = mp_hands.Hands(
    max_num_hands=1,
    min_detection_confidence=0.6,
    min_tracking_confidence=0.6,
    model_complexity=1
)

# 手势映射表
GESTURE_MAP = {0: "石头", 1: "剪刀", 2: "布", -1: "未识别", -2: "手背"}

# 从手部关键点提取特征
def extract_features(hand_landmarks):
    landmarks = hand_landmarks.landmark
    features = []

    # 计算关键点之间的相对距离和角度作为特征
    wrist = (landmarks[mp_hands.HandLandmark.WRIST].x, landmarks[mp_hands.HandLandmark.WRIST].y)

    for i in range(1, 21):  # 跳过手腕(0)
        point = (landmarks[i].x, landmarks[i].y)
        # 计算相对于手腕的位置
        features.append(point[0] - wrist[0])
        features.append(point[1] - wrist[1])

    # 计算手指之间的相对位置
    finger_tips = [mp_hands.HandLandmark.THUMB_TIP,
                   mp_hands.HandLandmark.INDEX_FINGER_TIP,
                   mp_hands.HandLandmark.MIDDLE_FINGER_TIP,
                   mp_hands.HandLandmark.RING_FINGER_TIP,
                   mp_hands.HandLandmark.PINKY_TIP]

    for i in range(len(finger_tips)):
        for j in range(i + 1, len(finger_tips)):
            x1, y1 = landmarks[finger_tips[i]].x, landmarks[finger_tips[i]].y
            x2, y2 = landmarks[finger_tips[j]].x, landmarks[finger_tips[j]].y
            # 计算距离
            dist = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
            features.append(dist)

    return features

# 数据收集函数
def collect_data():
    print("开始收集训练数据...")
    print("请按以下顺序做出手势：石头(0) -> 剪刀(1) -> 布(2)")
    print("每种手势收集100帧数据，按任意键开始...")
    cv2.waitKey(0)

    cap = cv2.VideoCapture(0)
    data = []
    labels = []

    for gesture_id in [0, 1, 2]:  # 石头、剪刀、布
        print(f"\n请做出 {GESTURE_MAP[gesture_id]} 的手势")
        print("3秒后开始收集数据...")
        time.sleep(3)

        count = 0
        while count < 100:  # 每种手势收集100帧
            success, frame = cap.read()
            if not success:
                print("无法获取图像")
                break

            frame = cv2.flip(frame, 1)
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = hands.process(rgb_frame)

            if results.multi_hand_landmarks:
                hand_landmarks = results.multi_hand_landmarks[0]
                # 绘制手部关键点
                mp.solutions.drawing_utils.draw_landmarks(
                    frame, hand_landmarks, mp_hands.HAND_CONNECTIONS)

                # 提取特征
                features = extract_features(hand_landmarks)
                data.append(features)
                labels.append(gesture_id)

                count += 1
                frame = cv2.putText(frame, f"收集进度: {count}/100", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            frame = cv2.putText(frame, f"请做出 {GESTURE_MAP[gesture_id]} 的手势", (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            cv2.imshow("收集训练数据", frame)

            if cv2.waitKey(1) == 27:  # 按ESC退出
                cap.release()
                cv2.destroyAllWindows()
                return None, None

        print(f"{GESTURE_MAP[gesture_id]} 手势数据收集完成")

    cap.release()
    cv2.destroyAllWindows()

    return np.array(data), np.array(labels)

# 训练决策树模型
def train_model(data, labels):
    print("开始训练模型...")
    model = DecisionTreeClassifier(max_depth=5, random_state=42)
    model.fit(data, labels)
    print("模型训练完成")

    # 保存模型
    with open('gesture_model.pkl', 'wb') as f:
        pickle.dump(model, f)
    print("模型已保存到 gesture_model.pkl")

    return model

# 加载预训练模型
def load_model():
    if os.path.exists('gesture_model.pkl'):
        try:
            with open('gesture_model.pkl', 'rb') as f:
                model = pickle.load(f)
            print("已加载预训练模型")
            return model
        except:
            print("无法加载预训练模型，将重新训练")
    return None

def main():
    # 尝试加载预训练模型
    model = load_model()

    # 如果没有预训练模型，收集数据并训练
    if model is None:
        data, labels = collect_data()
        if data is not None and len(data) > 0:
            model = train_model(data, labels)
        else:
            print("没有收集到训练数据，无法训练模型")

if __name__ == "__main__":
    main()